'''
Takes results from course alignment and pipe to exonerate to perform 
better alignment.
'''
import sys
import os
import subprocess

# Parameters.
quick_file = sys.argv[1]
est_file = sys.argv[2]
ref_file = sys.argv[3]
process_dir = sys.argv[4]
out_dir = sys.argv[5]
script_dir = sys.argv[6]


def read_file(filepath, split=False):
	# define data structure.
	data = {}
	
	# Check if zipped.
	zipped = False
	if filepath.count(".gz"):
		zipped = True
	
	# Open file.
	seq = ""
	seqheader = ""
	first = True
	
	if zipped == True:
		fin = gzip.open(filepath,"r")
	else:
		fin = open(filepath,"r")
	for line in fin:
		# Skip comments.	
		if line[0] == "#": continue

		# first time.
		if first == True and line[0] == ">":
			seq = ""
			seqheader = line.strip().replace(">","")
			first = False
			continue		

		# finished seq.
		if line[0] == ">":
			# save data.
			# See if we split header.
			if split == True:
				seqheader = seqheader.split()[0]
				
			data[seqheader] = seq.strip()

			# Reset variables.
			seq = ""
			seqheader = line.strip().replace(">","")
			continue

		# get seq.
		seq += line.strip()
			
	# Handle last close and return.
	if split == True:
		seqheader = seqheader.split()[0]
	data[seqheader] = seq.strip()
	fin.close()
	return data		
	


# Split up reference and EST.
estseqs = read_file(est_file, split=True)
refseqs = read_file(ref_file)

# Read in quick file.
genes = {}
fin = open(quick_file, "rb")
for line in fin:
	# Tokenize.
	tmp = line.strip().split()
	gene = tmp[0]
	scaf = tmp[1]
	
	# Add to library.
	if gene not in genes:
		genes[gene] = {}
	genes[gene][scaf] = True
	
# Make x number of lists.
cmds = []
for i in range(15):
	cmds.append([])
	
# Split up in several files.
i = 0
for gene in genes:
	# Make file names.
	out_reffile = "%s/ref_%s.fasta" % (process_dir, gene)
	out_estfile = "%s/est_%s.fasta" % (process_dir, gene)
	opath = "%s/%s.gff" % (out_dir, gene)
	
	# Create estfile.
	fout = open(out_estfile,"w")
	fout.write(">%s\n%s\n" % (gene, estseqs[gene]))
	fout.close()
	
	# Write reffile.
	fout = open(out_reffile,"w")
	for x in genes[gene]:
		fout.write(">%s\n%s\n" % (x, refseqs[x]))
	fout.close()
	
	# Build cmd.
	cmd = []
	cmd.append("exonerate")
	cmd.append("--showquerygff true")
	cmd.append("--showtargetgff true")
	cmd.append("--model est2genome")
	cmd.append(out_estfile)
	cmd.append(out_reffile)
	cmd.append("> %s" % (opath))
	cmd = ' '.join(cmd)
	cmds[i].append(cmd)
	
	# Switch i.
	i += 1
	if i >= 15:
		i = 0
	
# Split up commands over X files.
i = 0
for apple in cmds:
	# Write list.
	txt = "#!/bin/bash\n"
	txt += "\n".join(apple)
	txt += "\nexit;\n"

	# write script.
	spath = "%s/%i.sh" % (script_dir, i)
	sout = open(spath,"w")
	sout.write(txt)
	sout.close()
	
	# execute script.
	subprocess.call(["/home/jlindsay/batch_code/smp_sub.sh", spath])
	i += 1
	
	
fin.close()
